import numpy as np
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Activation
from tensorflow.keras.layers import LSTM
from sklearn.preprocessing import MinMaxScaler

import matplotlib.pyplot as plt

timesteps = seq_length = 7 # 时间步
data_dim = 5

# Open,High,Low,Close,Volume
xy = np.loadtxt('./data-02-stock_daily.csv', delimiter=',')
xy = xy[::-1]  # 数据倒叙，最后的数据，日期最远的数据

scaler = MinMaxScaler(feature_range=(0, 1))
xy = scaler.fit_transform(xy)

x = xy
y = xy[:, [-1]]  # Close as label

dataX = []
dataY = []
# 使用前7天预测第8天的闭盘价格
for i in range(0, len(y) - seq_length):
    _x = x[i:i + seq_length]
    _y = y[i + seq_length]
    print(_x, "->", _y)
    dataX.append(_x)
    dataY.append(_y)

# split to train and testing
train_size = int(len(dataY) * 0.7)
test_size = len(dataY) - train_size
trainX, testX = np.array(dataX[0:train_size]), np.array(
    dataX[train_size:len(dataX)])
trainY, testY = np.array(dataY[0:train_size]), np.array(
    dataY[train_size:len(dataY)])

model = Sequential()
model.add(LSTM(1, input_shape=(seq_length, data_dim), return_sequences=False)) # 只有最后返回预测结果
# model.add(Dense(1))
model.add(Activation("linear"))
model.compile(loss='mean_squared_error', optimizer='adam')

model.summary()

print(trainX.shape, trainY.shape)
model.fit(trainX, trainY, epochs=200)

# make predictions
testPredict = model.predict(testX)

plt.plot(testY)
plt.plot(testPredict)
plt.show()
